# -*- coding: utf-8 -*-
# @Time    : 2021/4/27 15:47
# @Author  : huangwei
# @File    : rec_system.py
# @Software: PyCharm
import copy
import os
import imutils
from config import args
from system_method import *
from word import TextDetector, TextRecognizer


class TextSystem(object):
    def __init__(self, args):
        self.text_detector = TextDetector(args)
        self.text_recognizer = TextRecognizer(args)

    def __call__(self, img):
        ori_im = img.copy()
        # 识别出有文字的框
        dt_boxes = self.text_detector(img)
        print("识别出{}个有文字的框".format(len(dt_boxes)))

        if dt_boxes is None:
            return None, None

        # 对box进行从上到下，从左到右的排序
        dt_boxes = sorted_boxes(dt_boxes)

        img_crop_list = []
        # 将识别出的框裁剪出来，再进行文字识别
        for bno in range(len(dt_boxes)):
            tmp_box = copy.deepcopy(dt_boxes[bno])
            img_crop = get_rotate_crop_image(ori_im, tmp_box)
            img_crop_list.append(img_crop)

        # 识别框中的文字
        rec_res = self.text_recognizer(img_crop_list)
        print(rec_res)

        # 返回识别率较高的框和结果
        filter_boxes, filter_rec_res = [], []
        total_result = []

        # 返回框和识别的结果，一一对应
        for box, rec_reuslt in zip(dt_boxes, rec_res):
            line = {}
            text, score = rec_reuslt
            if score >= 0.5:
                filter_boxes.append(box)
                filter_rec_res.append(rec_reuslt)
                line["box"] = box
                line["rec"] = rec_reuslt
                total_result.append(line)
        return filter_boxes, filter_rec_res, total_result


def get_data(img):
    # 检测图片中的文字框
    dt_boxes, _, _ = text_sys(img)

    # 根据文字框计算出需要旋转的角度
    roll_angle = get_angle(dt_boxes)
    roll_img = imutils.rotate_bound(img, -roll_angle)

    # 找到写有站名的两个box判断其是否在图片的上半区域，否则旋转180°
    dt_boxes, rec_res, result = text_sys(roll_img)
    shape = roll_img.shape[:2]
    need_flip = need_to_flip(result, shape)
    if need_flip:
        roll_img = imutils.rotate_bound(roll_img, 180)
        dt_boxes, rec_res, result = text_sys(roll_img)

    # 现在已经正放了，找出识别文字中需要的信息
    shape = roll_img.shape[:2]
    all_infos = get_infomations(result, shape)

    return all_infos


text_sys = TextSystem(args)

if __name__ == '__main__':
    try:
        img_path = "pic/code3.png"
        img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif'}
        if os.path.isfile(img_path) and os.path.splitext(img_path)[-1][1:].lower() in img_end:
            img = cv2.imread(img_path)
            # 判断图片是否过大，是的话则转为长宽最大为2000的图片
            img = normal_size(img)

            infos = get_data(img)
            print(infos)
        else:
            print("image is not exist!")
    except:
        pass
